from pykalman import KalmanFilter as KF
from random import gauss
import numpy as np
from colorama import init, Fore
import matplotlib.pyplot as plt


init(autoreset=True)

A = np.array([[0.98, -0.08, 0.99], [0.04, 1.01, 0.0],[0.01,0.0,0.99]])
C = np.array([[1.0, 0.0, 0.01], [0.0, 1.0, 0.0],[0.0,0.0,1.0]])

def format_array(arr):
	result = ""
	for row in arr:
		result += Fore.BLUE + " ".join([str(x) for x in row]) + "\n"
	return result

def gauss_noise(mu, sigma, dim):
	return np.array([[gauss(mu, sigma)] for i in range(dim)])

states = [np.array([[1.0], [1.0], [1.0]])]
obs = [np.dot(C, states[0]) + gauss_noise(0.0,0.2,3)]

for i in range(19):
	states.append(np.dot(A, states[i]) + gauss_noise(0.0, 0.1, 3))
	obs.append(np.dot(C, states[i+1]) + gauss_noise(0.0, 0.2, 3))
"""
for i, state in enumerate(states):
	print "------------------------------------------"
	print "STATE:\n", format_array(state)
	print "OBSERVATION:\n", format_array(obs[i])
"""

kf = KF(initial_state_mean=[1.0,1.0,0.0,0.0],  n_dim_obs=3, n_dim_state=4,
        em_vars = 'all')

kf = kf.em([[ob[0][0], ob[1][0], ob[2][0]] for ob in obs], n_iter=100)

print "OBS:\n", format_array(kf.observation_matrices)
print "STATE: \n", format_array(kf.transition_matrices)


smooth_estimates = kf.smooth([[ob[0][0], ob[1][0], ob[2][0]] for ob in obs])[0]

filtered_estimates = kf.filter([[ob[0][0], ob[1][0], ob[2][0]] for ob in obs])[0]


x_axis = np.arange(20)

plt.plot(x_axis, [x[0] for x in states],"g", x_axis, [x[0] for x in obs], "g--",
        x_axis,[x[0] for x in  smooth_estimates], "b", x_axis, [x[0] for x in
            filtered_estimates], "b--",  x_axis, [x[0] for x in
            kf.sample(20)[1]], "r")
plt.show()
plt.savefig("output.png")
#plt.plot(x_axis, [x[1] for x in states],"g", x_axis, [x[1] for x in obs], "r", x_axis,[x[1] for x in  kf.sample(100)[1]], "b")
#plt.show()




